# 使用transformer原生库调用

from modelscope import AutoModelForCausalLM, AutoTokenizer

model_name = '/home/aresen/1project/2python/hub/DeepSeek-R1-Distill-Qwen-1.5B'

# 实例化预训练模型和分词器
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype='auto',
    device_map='auto',
    low_cpu_mem_usage=True
)

tokenizer = AutoTokenizer.from_pretrained(model_name)

# 创建消息message

prompt = "你好，好久不见，请介绍下自己"
message = [
    {'role': 'system', 'content': '你是一名助人为乐的助手'},
    {'role': 'user', 'content': prompt}
]

# 词嵌入过程
text = tokenizer.apply_chat_template(
    message,
    tokenize=False,
    add_generation_prompt=True
)

model_inputs = tokenizer([text], return_tensors='pt').to(model.device)

# 创建并回复
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=512
)
generated_ids=[
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

print(response)

